import math
import torch
import torch.nn as nn
from torch.nn.parameter import Parameter
from . import block as B
import torch.nn.functional as F


class Scale(nn.Module):

    def __init__(self, init_value=1e-3):
        super().__init__()
        self.scale = nn.Parameter(torch.FloatTensor([init_value]))

    def forward(self, input):
        return input * self.scale


class AWRU(nn.Module):
    def __init__(
        self, n_feats, kernel_size, block_feats, wn, res_scale=1, act=nn.ReLU(True)):
        super(AWRU, self).__init__()
        self.res_scale = Scale(1)
        self.x_scale = Scale(1)
        body = []
        body.append(
            wn(nn.Conv2d(n_feats, block_feats, kernel_size, padding=kernel_size//2)))
        body.append(act)
        body.append(
            wn(nn.Conv2d(block_feats, n_feats, kernel_size, padding=kernel_size//2)))

        self.body = nn.Sequential(*body)

    def forward(self, x):
        res = self.res_scale(self.body(x)) + self.x_scale(x)
        return res


class AWMS(nn.Module):
    def __init__(
        self, n_colors, scale, n_feats, kernel_size, wn):
        super(AWMS, self).__init__()
        out_feats = scale*scale*n_colors
        self.tail_k3 = wn(nn.Conv2d(n_feats, out_feats, 3, padding=3//2, dilation=1))
        self.tail_k5 = wn(nn.Conv2d(n_feats, out_feats, 5, padding=5//2, dilation=1))
        self.pixelshuffle = nn.PixelShuffle(scale)
        self.scale_k3 = Scale(0.5)
        self.scale_k5 = Scale(0.5)


    def forward(self, x):
        #x0 = self.pixelshuffle(self.scale_k3(self.tail_k3(x)))
       # x1 = self.pixelshuffle(self.scale_k5(self.tail_k5(x)))
        x0 = self.pixelshuffle(self.scale_k3(self.tail_k3(x)))
        x1 = self.pixelshuffle(self.scale_k5(self.tail_k5(x)))

        return x0+x1#+x2+x3

def pixel_unshuffle(input, downscale_factor):
    '''
    input: batchSize * c * k*w * k*h
    kdownscale_factor: k
    batchSize * c * k*w * k*h -> batchSize * k*k*c * w * h
    '''
    c = input.shape[1]

    kernel = torch.zeros(size=[downscale_factor * downscale_factor * c,
                               1, downscale_factor, downscale_factor],
                         device=input.device)
    for y in range(downscale_factor):
        for x in range(downscale_factor):
            kernel[x + y * downscale_factor::downscale_factor*downscale_factor, 0, y, x] = 1
    return F.conv2d(input, kernel, stride=downscale_factor, groups=c)

class PixelUnshuffle(nn.Module):
    def __init__(self, downscale_factor):
        super(PixelUnshuffle, self).__init__()
        self.downscale_factor = downscale_factor
    def forward(self, input):
        '''
        input: batchSize * c * k*w * k*h
        kdownscale_factor: k
        batchSize * c * k*w * k*h -> batchSize * k*k*c * w * h
        '''

        return pixel_unshuffle(input, self.downscale_factor)


def channel_shuffle(x, groups):
    batchsize, num_channels, height, width = x.data.size()

    channels_per_group = num_channels // groups

    # reshape
    x = x.view(batchsize, groups,
               channels_per_group, height, width)

    # transpose
    # - contiguous() required if transpose() is used before view().
    #   See https://github.com/pytorch/pytorch/issues/764
    x = torch.transpose(x, 1, 2).contiguous()

    # flatten
    x = x.view(batchsize, -1, height, width)

    return x
class LFB(nn.Module):
    def __init__(
        self, n_feats, wn, act=nn.ReLU(True)):
        super(LFB, self).__init__()

      #  n_feats = n_feats//4
        self.b0 = B.ESAB_MODIFIED(n_feats=n_feats,reduction_factor=4)#AWRU(n_feats, kernel_size, block_feats, wn=wn, act=act)
        self.b1 = B.ESAB_MODIFIED(n_feats=n_feats,reduction_factor=4)#AWRU(n_feats, kernel_size, block_feats, wn=wn, act=act)
        self.b2 = B.ESAB_MODIFIED(n_feats=n_feats,reduction_factor=4)#AWRU(n_feats, kernel_size, block_feats, wn=wn, act=act)
        self.b3 = B.ESAB_MODIFIED(n_feats=n_feats,reduction_factor=4)#AWRU(n_feats, kernel_size, block_feats, wn=wn, act=act)
        #self.b4 = B.ESAB_MODIFIED(n_feats=n_feats,reduction_factor=4)  # AWRU(n_feats, kernel_size, block_feats, wn=wn, act=act)
        #self.b5 = B.ESAB(n_feats=n_feats)# AWRU(n_feats, kernel_size, block_feats, wn=wn, act=act)

     #   self.conv = nn.Conv2d(n_feats,n_feats,3,stride=2,padding=1)
        self.reduction1 = wn(nn.Conv2d(n_feats*2, n_feats, 1))
        self.reduction2 = wn(nn.Conv2d(n_feats*2, n_feats, 1))
        self.reduction3 = wn(nn.Conv2d(n_feats*2, n_feats, 1))
        self.res_scale = Scale(1)
        self.x_scale = Scale(1)

       # self.up   = nn.PixelShuffle(2)
        #self.down = PixelUnshuffle(2)

    def forward(self, x):

      #  x = self.up(x)
      #  print(x.shape)
        x0 = self.b0(x)
        x1 = self.b1(x0)+x0
        x2 = self.b2(x1)+x1
        x3 = self.b3(x2)
      #  x4 = self.b4(x3)
     #   x5 = self.b5(x4)
        res1 = self.reduction1(channel_shuffle(torch.cat([x0, x1],dim=1), 2))
        res2 = self.reduction2(channel_shuffle(torch.cat([res1, x2], dim=1), 2))
        res = self.reduction3(channel_shuffle(torch.cat([res2,x3], dim=1), 2))


        #res = self.down(res)
        return self.res_scale(res) + self.x_scale(x)

class IMDN(nn.Module):
    def __init__(self, upscale=4):
        super(IMDN, self).__init__()
        # hyper-params
        ########################################################################

        scale = 4
        n_resblocks = 4
        n_feats = 32
        kernel_size = 3
        act = nn.LeakyReLU(True)
        n_colors= 3

        # wn = lambda x: x
        wn = lambda x: torch.nn.utils.weight_norm(x)

        self.rgb_mean = torch.autograd.Variable(torch.FloatTensor(
            [0.4488, 0.4371, 0.4040])).view([1, 3, 1, 1])

        # define head module
        # head = HEAD(args, n_feats, kernel_size, wn)
        head = []
        head.append(
            wn(nn.Conv2d(3, n_feats, 3, padding=3//2)))

        # define body module
        body = []
        for i in range(n_resblocks):
            body.append(
                LFB(n_feats, wn=wn, act=act))

        # define tail module
       # out_feats = scale*scale*n_colors
        tail = AWMS(n_colors, scale, n_feats, kernel_size, wn)

        # skip = []
        # skip.append(
        #     wn(nn.Conv2d(n_colors, out_feats, 3, padding=3//2))
        # )
        # skip.append(nn.PixelShuffle(scale))

        # make object members
        self.head = nn.Sequential(*head)
        self.body = nn.Sequential(*body)
        self.tail = tail
       # self.skip = nn.Sequential(*skip)


    def forward(self, x):
        input = x
        #x = (x - self.rgb_mean.cuda()*255)/127.5
       # s = self.skip(x)
        x = self.head(x)
        x = self.body(x)
        x = self.tail(x)
      #  x += s
        #x = x*127.5 + self.rgb_mean.cuda()*255
        return x + torch.nn.functional.upsample(input, scale_factor=4, mode='bicubic')


    def load_state_dict(self, state_dict, strict=True):
        own_state = self.state_dict()
        for name, param in state_dict.items():
            if name in own_state:
                if isinstance(param, nn.Parameter):
                    param = param.data
                try:
                    own_state[name].copy_(param)
                except Exception:
                    if name.find('tail') >= 0 or  name.find('skip') >= 0:
                        print('Replace pre-trained upsampler to new one...')
                    else:
                        raise RuntimeError('While copying the parameter named {}, '
                                           'whose dimensions in the model are {} and '
                                           'whose dimensions in the checkpoint are {}.'
                                           .format(name, own_state[name].size(), param.size()))
            elif strict:
                if name.find('tail') == -1:
                    raise KeyError('unexpected key "{}" in state_dict'
                                   .format(name))

        if strict:
            missing = set(own_state.keys()) - set(state_dict.keys())
            if len(missing) > 0:
                raise KeyError('missing keys in state_dict: "{}"'.format(missing))